examples/kft-v2/manifests/shared_pvc.yaml (1)

12-12: Add missing trailing newline.

Fix YAML lint error (new-line-at-end-of-file).

-  volumeMode: Filesystem
+  volumeMode: Filesystem
+

examples/kft-v2/manifests/cluster_training_runtime.yaml (2)

116-119: Persist progression file on the shared volume.

Writing to /tmp is ephemeral; put the progression JSON on /workspace so it persists and is visible to controllers.

-                        - name: TRAINJOB_PROGRESSION_FILE_PATH
-                          value: /tmp/training_progression.json
+                        - name: TRAINJOB_PROGRESSION_FILE_PATH
+                          value: /workspace/training_progression.json

---

106-115: NCCL env on CPU-only runs is unnecessary.

If this runtime is CPU-only, consider dropping NCCL_* envs to simplify. The script already selects gloo on CPU.

examples/kft-v2/scripts/mnist.py (5)

58-61: Fix type hints to use Optional[...] per PEP 484.

Adjust signatures to satisfy Ruff RUF013.

-            loss: float = None,
-            learning_rate: float = None,
-            checkpoint_dir: str = None,
+            loss: Optional[float] = None,
+            learning_rate: Optional[float] = None,
+            checkpoint_dir: Optional[str] = None,
@@
-        def update_epoch(self, epoch: int, checkpoint_dir: str = None, **metrics):
+        def update_epoch(self, epoch: int, checkpoint_dir: Optional[str] = None, **metrics):

Also applies to: 127-127

---

1-1: Remove shebang or mark file executable.

To silence EXE001, either remove the shebang (common for importable modules) or set execute bit. Removing is simplest.

-#!/usr/bin/env python3

---

45-46: Default progression file should live on a mounted volume, not /tmp.

Align default with the manifest so status persists and is observable.

-            self.status_file_path = status_file_path or os.getenv(
-                "TRAINJOB_PROGRESSION_FILE_PATH", "/tmp/training_progression.json"
+            self.status_file_path = status_file_path or os.getenv(
+                "TRAINJOB_PROGRESSION_FILE_PATH", "/workspace/training_progression.json"

---

471-471: Avoid modifying grads when reporting grad norm.

clip_grad_norm_ changes gradients; compute norm without clipping or move clipping before optimizer.step if intended.

Add a helper and use it here:

def compute_grad_norm(parameters, norm_type=2.0):
    params = [p for p in parameters if p.grad is not None]
    if not params:
        return 0.0
    if norm_type == float("inf"):
        return max(p.grad.detach().abs().max().item() for p in params)
    total = 0.0
    for p in params:
        param_norm = p.grad.detach().data.norm(norm_type)
        total += param_norm.item() ** norm_type
    return total ** (1.0 / norm_type)

-                        grad_norm=f"{torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0):.4f}"
+                        grad_norm=f"{compute_grad_norm(model.parameters()):.4f}"

---

224-225: Narrow overly broad exception handlers.

Catching bare Exception hides actionable errors. Scope to expected exceptions (e.g., RuntimeError for DDP init; OSError/ValueError for file ops).

Also applies to: 285-286, 340-341

examples/kft-v2/scripts/trl_training.py (7)

556-558: Remove unused variable.

checkpoint_interval is never used.

-        checkpoint_interval = os.getenv('CHECKPOINT_INTERVAL', '30s')

---

606-606: Drop redundant import.

os is already imported at Line 4.

-    import os

---

272-274: Don’t swallow exceptions silently.

Replace bare except/pass with minimal logging.

-            except Exception as e:
-                pass
+            except Exception as e:
+                print(f"Failed to write simple progress: {e}")

---

65-65: Consistent step indexing (+1).

mnist.ProgressionTracker uses 1-based steps; mirror for consistency.

-                self.current_step = (epoch - 1) * self.steps_per_epoch + step
+                self.current_step = (epoch - 1) * self.steps_per_epoch + step + 1

---

51-52: Progress file placement and durability.

Defaulting to /tmp risks loss on pod restart. Prefer a PVC path like CHECKPOINT_URI or a dedicated PROGRESSION_URI on RWX storage.

Option: set TRAINJOB_PROGRESSION_FILE_PATH to f"{CHECKPOINT_URI}/training_progression.json" and ensure directory exists before writing.

Also applies to: 181-185, 270-271, 774-782

---

663-669: Minor: remove f-prefix on constant strings.

-                print(f"Applied DDP static graph fix for distributed training")
+                print("Applied DDP static graph fix for distributed training")
@@
-                    print(f"PEFT model detected, DDP parameters properly configured")
+                    print("PEFT model detected, DDP parameters properly configured")

---

37-187: Deduplicate utilities.

ProgressionTracker and setup_distributed duplicate mnist.py versions. Factor into a shared helper (e.g., examples/kft-v2/scripts/common.py) for reuse.

I can extract and wire a common module if helpful.

Also applies to: 424-454

examples/kft-v2/trl-gpt2-checkpointing.ipynb (3)

624-626: Progress file path on PVC instead of /tmp.

Use a RWX path so UIs can read it across restarts.

Example:

-    "        checkpoint_dir = Path(os.getenv('CHECKPOINT_URI', '/workspace/checkpoints'))\n",
+    "        checkpoint_dir = Path(os.getenv('CHECKPOINT_URI', '/workspace/checkpoints'))\n",
+    "        os.environ.setdefault('TRAINJOB_PROGRESSION_FILE_PATH', str(checkpoint_dir / 'training_progression.json'))\n",

And in env:

-    "    \"TRAINJOB_PROGRESSION_FILE_PATH\": \"/tmp/training_progression.json\",\n",
+    "    \"TRAINJOB_PROGRESSION_FILE_PATH\": \"/workspace/checkpoints/training_progression.json\",\n",

Also applies to: 916-917, 833-833

---

768-818: Reuse shared utilities to avoid drift.

ProgressionTracker and setup_distributed duplicate the script and mnist.py. Extract to a common module.

I can factor these into examples/kft-v2/scripts/common.py and update imports.

---

955-955: Minor: remove f-strings without placeholders.

-    "                print(f\"Applied DDP static graph fix for distributed training\")\n",
+    "                print(\"Applied DDP static graph fix for distributed training\")\n",
@@
-    "                    print(f\"PEFT model detected, DDP parameters properly configured\")\n",
+    "                    print(\"PEFT model detected, DDP parameters properly configured\")\n",
@@
-    "    \"print(f\\\"Trainjob submitted!!\\\")\"\n",
+    "    \"print(\\\"Trainjob submitted!!\\\")\"\n",

Also applies to: 963-963, 735-741

Configuration used: CodeRabbit UI

Review profile: CHILL

Plan: Pro

examples/kft-v2/scripts/trl_training.py (4)

10-11: Critical security issue flagged in previous review remains unaddressed.

The code still imports numpy.core.multiarray._reconstruct, registers it as a safe global (line 35), and patches torch.load to default weights_only=False (lines 189-195). This approach widens the RCE attack surface during model deserialization.

Based on previous review feedback, change to safer defaults.

Also applies to: 35-35, 189-195

---

442-451: Past issue remains: hard-coded Gloo backend loses GPU performance.

The code still hard-codes backend='gloo' (line 445), which was flagged in previous reviews. For GPU training, NCCL offers significantly better performance. See mnist.py lines 267-268 for the correct pattern:

device, backend = ("cuda", "nccl") if torch.cuda.is_available() else ("cpu", "gloo")

---

466-468: Past issue remains: negative split size when dataset < 20 samples.

Lines 466-468 can produce negative train_size when len(full_dataset) < 20. This was flagged in previous reviews.

Apply the suggested fix:

-                    train_size = min(100, len(full_dataset) - 20)
-                    train_dataset = full_dataset.select(range(train_size))
-                    test_dataset = full_dataset.select(range(train_size, min(train_size + 20, len(full_dataset))))
+                    total = len(full_dataset)
+                    test_sz = min(20, max(0, total))
+                    train_sz = min(100, max(0, total - test_sz))
+                    train_dataset = full_dataset.select(range(train_sz))
+                    test_dataset = full_dataset.select(range(train_sz, min(train_sz + test_sz, total)))

---

577-578: Past issues remain: incorrect precision flags and hard-coded backend.

Two problems flagged in previous reviews:

1. Lines 577-578: fp16 on CPU will crash; bf16 needs GPU support check
2. Line 596: ddp_backend='gloo' hard-coded, losing GPU performance

Apply the precision fix:

-            'bf16': torch.cuda.is_available(),  # Only use bf16 if CUDA is available
-            'fp16': not torch.cuda.is_available(),  # Use fp16 for CPU training
+            'bf16': torch.cuda.is_available() and bool(getattr(torch.cuda, 'is_bf16_supported', lambda: False)()),
+            'fp16': False,  # Disable fp16 unless explicitly needed

And the backend fix:

-            'ddp_backend': 'gloo',  # Use gloo backend for better LoRA compatibility
+            'ddp_backend': 'nccl' if torch.cuda.is_available() else 'gloo',

Also applies to: 596-596

examples/kft-v2/scripts/mnist.py (4)

292-544: Critical: Nested function never called, training is a no-op.

The outer train_fashion_mnist() (line 2) defines an inner train_fashion_mnist() (line 292) but never calls it. When main() invokes the outer function, no training executes. This was flagged in previous reviews.

Rename the inner function and call it:

-    def train_fashion_mnist():
+    def _run_training():
         # Setup distributed training
         local_rank, global_rank, world_size, device_type = setup_distributed()
         ...
         # Finally clean up PyTorch distributed
         if world_size > 1:
             dist.destroy_process_group()
+    
+    # Run the actual training
+    _run_training()

---

176-179: Past issue remains: off-by-one in epoch display.

Line 176 adds 1 to epoch, but the docstring (line 65) and usage (line 516) indicate epoch is already the absolute 1-based epoch number, causing the display to show epoch+1.

Apply this fix:

-            epoch_num = epoch + 1
             total_epochs = self.total_epochs
-            message = f"Completed epoch {epoch_num}/{total_epochs}"
+            message = f"Completed epoch {epoch}/{total_epochs}"

---

304-307: Past issue remains: DDP missing device_ids for multi-GPU nodes.

Line 305 wraps the model without specifying device_ids and output_device, which can cause warnings on multi-GPU nodes. This was flagged in previous reviews.

Apply this fix:

         # Create model and wrap with DDP only if distributed
         net = Net().to(device)
         if world_size > 1:
-            model = nn.parallel.DistributedDataParallel(net)
+            if device.type == "cuda":
+                model = nn.parallel.DistributedDataParallel(
+                    net, device_ids=[local_rank], output_device=local_rank
+                )
+            else:
+                model = nn.parallel.DistributedDataParallel(net)
         else:
             model = net

---

467-467: Past issue remains: dist.get_world_size() raises when not initialized.

Line 467 calls dist.get_world_size() which will raise in single-process mode. Use the world_size variable from setup_distributed() instead, as flagged in previous reviews.

Apply this fix:

-                        world_size=dist.get_world_size(),
+                        world_size=world_size,

examples/kft-v2/trl-gpt2-checkpointing.ipynb (3)

123-169: Past issue remains: duplicate environment variable keys.

Lines 123, 155 duplicate PYTHONUNBUFFERED; lines 124, 161 duplicate NCCL_DEBUG; lines 125, 167 duplicate TORCH_DISTRIBUTED_DEBUG. Later values overwrite earlier ones, as flagged in previous reviews.

Remove the duplicates:

     "PYTHONUNBUFFERED": "1",
     "NCCL_DEBUG": "INFO",
     "TORCH_DISTRIBUTED_DEBUG": "INFO",
     "PYTHONPATH": "/tmp/lib:$PYTHONPATH",
     ...
     # Cache directories
-    "PYTHONUNBUFFERED": "1",
     "TRANSFORMERS_CACHE": "/workspace/cache/transformers",
     ...
     # Distributed training debug
-    "NCCL_DEBUG": "INFO",
     "NCCL_DEBUG_SUBSYS": "ALL",
     ...
-    "TORCH_DISTRIBUTED_DEBUG": "INFO",

---

183-193: Past issue remains: pip flags in packages list and missing commas.

Lines 191-192 include pip flags as list items, and line 190 is missing a comma, causing string concatenation. This was flagged in previous reviews.

Fix by removing pip flags and adding comma:

     packages_to_install=[
         "transformers[torch]",
         "trl", 
         "peft", 
         "datasets", 
         "accelerate",
         "torch",
-        "numpy"
-        " --target=/tmp/lib"
-        " --verbose"
+        "numpy",
     ],

If you need to pass additional pip arguments, check if Kubeflow SDK supports a dedicated parameter for pip flags.

---

244-252: Past security issues remain: hardcoded credentials and disabled TLS.

Lines 244-245 use placeholder strings for credentials that users might accidentally commit, and line 252 disables SSL verification. These were flagged in previous reviews.

Use environment variables:

+import os
-api_server = "<api-server-url>"
-token = "<auth-token>"
+api_server = os.getenv("KFP_API_SERVER", "")
+token = os.getenv("KFP_API_TOKEN", "")

 configuration = client.Configuration()
 configuration.host = api_server
 configuration.api_key = {"authorization": f"Bearer {token}"}

-# Un-comment if your cluster API server uses a self-signed certificate or an un-trusted CA
-configuration.verify_ssl = False
+# Only disable SSL verification if explicitly needed
+if os.getenv("KFP_VERIFY_SSL", "true").lower() == "false":
+    configuration.verify_ssl = False

examples/kft-v2/scripts/trl_training.py (3)

37-188: Extract ProgressionTracker to eliminate duplication.

ProgressionTracker is nearly identical in trl_training.py (lines 37-188) and mnist.py (lines 21-225), with only minor differences in handling global_step and ETA formatting. Extract to a shared utility module to improve maintainability.

---

556-556: Remove unused variable.

Line 556: checkpoint_interval is assigned but never used.

Apply this diff:

         checkpoint_dir = Path(os.getenv('CHECKPOINT_URI', '/workspace/checkpoints'))
         checkpoint_enabled = os.getenv('CHECKPOINT_ENABLED', 'false').lower() == 'true'
-        checkpoint_interval = os.getenv('CHECKPOINT_INTERVAL', '30s')
         max_checkpoints = int(os.getenv('CHECKPOINT_MAX_RETAIN', '5'))

---

606-606: Remove duplicate import.

Line 606 re-imports os, which was already imported at line 4.

Apply this diff:

     """Training function."""
-
-    import os

examples/kft-v2/scripts/mnist.py (1)

1-1: Make file executable if shebang is intended.

Line 1 has a shebang but the file is not executable. Either make it executable with chmod +x or remove the shebang if not needed.

examples/kft-v2/trl-gpt2-checkpointing.ipynb (1)

171-171: Add import clarification for better UX.

Line 171 imports trl_train from trl_training.py. Consider adding a cell comment explaining that users need the scripts/trl_training.py file in their Python path, or show how to add it:

import sys
sys.path.append('./scripts')
from trl_training import trl_train

Configuration used: CodeRabbit UI

Review profile: CHILL

Plan: Pro